#ifndef DSLKJLE_20171106_32LSKDFE_
#define DSLKJLE_20171106_32LSKDFE_

#include "define.hpp"
using namespace PiaCNN;

#define POOLING_MAX		 0
#define POOLING_MEAN	 1

class PiaPool
{
public:
	PiaPool();
	~PiaPool();

public:
	bool Create (
				int nInputRows, 
				int nInputCols, 
				int nInputChan,
				int filterSize1, 
				int filterSize2,
				int stride1, 
				int stride2,
				int PoolType = POOLING_MAX,
			 string name = "");
	bool Run(std::vector<MatrixXd> &input, 
			 std::vector<MatrixXd> &dst);
	bool Run(PiaTensor &srcOri,
			 PiaTensor &dstOri);
	void Print();
	void PrintResult(const std::vector<MatrixXd> &dst);
private:
	bool PreparePoolInput(const MatrixXd &input);
	
private:
	int dbg;
	
	
public:
	int nIdx;
	string name;
	
	int nInputRows, nInputCols, nInputChan;
	int filterSize1, filterSize2, filterNumber;
	int stride1, stride2;
	int PoolType;
	int nTop, nBottom, nLeft, nRight;
	int paddedRows, paddedCols;

	int nOutputRows, nOutputCols, nOutputChan;

private:
	MatrixXd prepareMultMap;
	MatrixXd OutputMaps;
};

PiaPool::PiaPool()
{
	nIdx		= -1;
	name		= "";
	nInputRows 	= 0;
	nInputCols 	= 0;
	nInputChan 	= 0;
	filterSize1 = 0;
	filterSize2 = 0;
	filterNumber= 0; // no need in pooling
	stride1 	= 0;
	stride2 	= 0;
	PoolType	= 0;
	dbg 		= 0;
}

PiaPool::~PiaPool()
{

}

/************************************************
* Calculate the padding value of the four directories
* and the size of padded map.
*/
bool PiaPool::Create (
						int nInputRows, 
						int nInputCols, 
						int nInputChan,
						int filterSize1, 
						int filterSize2,
						int stride1, 
						int stride2,
						int PoolType,
					 string name)
{
	if (nInputRows < 3 || nInputCols < 3, nInputChan < 1) {
		cout << ">> error. input data size" << endl;
		return false;
	}
	
	if (filterSize1 < 1 || filterSize2 < 1) {
		cout << ">> error. The smallest mask of pooling is 1*1" << endl;
		return false;
	}
	
	if (stride1 < 1 || stride2 < 1) {
		cout << ">> error. The smallest stride is 1" << endl;
		return false;
	}
	
	if (stride1 > filterSize1 || stride2 > filterSize2) {
		cout << ">> error. stride must <= filter size" << endl;
		return false;
	}
	if (PoolType != POOLING_MAX && PoolType != POOLING_MEAN) {
		cout << ">> error. pooling type" << endl;
		return false;
	}
	this->name 		  = name;
	this->nInputRows  = nInputRows;
	this->nInputCols  = nInputCols;
	this->nInputChan  = nInputChan;
	this->filterSize1 = filterSize1;
	this->filterSize2 = filterSize2;
	this->stride1     = stride1;
	this->stride2     = stride2;
	this->PoolType	  = PoolType;
	
	int nH = filterSize1 - 1;
	int nW = filterSize2 - 1;
	
	nTop = nH / 2;
	nBottom = nTop;
	if (nH % 2 != 0) 
	{
		nBottom += 1;
	}
	
	nLeft = nW / 2;
	nRight = nLeft;
	if (nW % 2 != 0) 
	{
		nRight += 1;
	}
	paddedRows = nInputRows + nTop + nBottom;
	paddedCols = nInputCols + nLeft + nRight;
	
	nOutputRows = (nInputRows - 1) / stride1 + 1;
	nOutputCols = (nInputCols - 1) / stride2 + 1;
	nOutputChan = nInputChan;
	
	prepareMultMap = MatrixXd(nInputRows * nInputCols, filterSize1 * filterSize2);
	OutputMaps = MatrixXd(nOutputRows * nOutputCols, nInputChan);
	
	if (dbg == 1) Print();
	return true;
}

bool PiaPool::PreparePoolInput(const MatrixXd &input)
{
	// padding
	MatrixXd padded;
	switch (1)  // padding method
	{
		case 0: // 1, 2, 3, 0, -1, 1 |  0,  0 
		{
			padded = MatrixXd::Zero(paddedRows, paddedCols);
			break;
		}
		case 1:	// 1, 2, 3, 0, -1, 1 | -1, -1
		{
			padded = MatrixXd(paddedRows, paddedCols);
			padded.setConstant(input.minCoeff());
			break;
		}
	}
	padded.block(nTop, nLeft, nInputRows, nInputCols) = input;
	
	// reshape
	for (int n2 = 0; n2 < filterSize2; n2++)
	{
		for (int n1 = 0; n1 < filterSize1; n1++)
		{
			for (int j = 0; j < nInputCols; j += 1)
			{
				for (int i = 0; i < nInputRows; i += 1)
				{
					prepareMultMap(j * nInputRows + i, n2 * filterSize1 + n1) = 
					padded(i + n1, j + n2);
				}
			}
		}
	}
	return true;
}

bool PiaPool::Run(PiaTensor &srcOri, PiaTensor &dstOri)
{
	PiaConv2dData *src = (struct PiaConv2dData*)&srcOri;
	PiaConv2dData *dst = (struct PiaConv2dData*)&dstOri;
	
	if (src->IsValid() == false) {
		cout << ">> error. pool.hpp Run() IsValid()" << endl;
		return false;
	}
	for (int i = 0; i < src->data.cols(); i++) {
		PreparePoolInput(MatrixXd::Map(src->data.middleCols(i, 1).data(), 
		                               src->shape[1], 
		                               src->shape[2]));
		MatrixXd TMP;
		switch (PoolType)
		{
			case POOLING_MAX: {
				TMP = prepareMultMap.rowwise().maxCoeff();
				break;
			}
			case POOLING_MEAN: {
				TMP = prepareMultMap.rowwise().mean();
			}
			default: {
				TMP = prepareMultMap.rowwise().maxCoeff();
				break;
			}
		}
		// stride 
		for (int c = 0; c < nOutputCols; c += 1) {
			for (int r = 0; r < nOutputRows; r += 1) {
				int nsrcIdx = c * nInputRows * stride2 + r * stride1;
				int ndstIdx = c * nOutputRows + r;
				OutputMaps(ndstIdx, i) = TMP(nsrcIdx, 0);
			}
		}
	}
	
	dst->data  = OutputMaps;
	dst->shape[1] = nOutputRows;
	dst->shape[2] = nOutputCols;
	dst->shape[3] = nOutputChan; 
	return true;
}

/************************************************
* Max pooling
*/
bool PiaPool::Run(std::vector<MatrixXd> &input, std::vector<MatrixXd> &dst)
{
	if ((int)input.size() != nInputChan) {
		cout << ">> error. input channels != model channels" << endl;
		return false;
	}
	
	for (int i = 0; i < nInputChan; i++) {
		PreparePoolInput(input[i]);
		MatrixXd TMP;
		switch (PoolType)
		{
			case POOLING_MAX: {
				TMP = prepareMultMap.rowwise().maxCoeff();
				break;
			}
			case POOLING_MEAN: {
				TMP = prepareMultMap.rowwise().mean();
			}
			default: {
				TMP = prepareMultMap.rowwise().maxCoeff();
				break;
			}
		}
		// stride 
		for (int c = 0; c < nOutputCols; c += 1) {
			for (int r = 0; r < nOutputRows; r += 1) {
				int nsrcIdx = c * nInputRows * stride2 + r * stride1;
				int ndstIdx = c * nOutputRows + r;
				OutputMaps(ndstIdx, i) = TMP(nsrcIdx, 0);
			}
		}
	}
	
	// reshape the output data
	dst.clear();
	for (int i = 0; i < OutputMaps.cols(); i++) {
		MatrixXd curFM= MatrixXd::Map(&OutputMaps.col(i)[0], nOutputRows, nOutputCols);
		dst.push_back(curFM);
	}
	
	if (dbg == 1) PrintResult(dst);
	return true;
}

void PiaPool::Print()
{
	int n = 9;
	cout << ">>                  " << endl;
	cout << ">> nInputRows......:" << setw(n) << nInputRows << endl;
	cout << ">> nInputCols......:" << setw(n) << nInputCols << endl;
	cout << ">> nInputChan......:" << setw(n) << nInputChan << endl;
	cout << ">> filterSize1.....:" << setw(n) << filterSize1 << endl;
	cout << ">> filterSize2.....:" << setw(n) << filterSize2 << endl;
	cout << ">> stride1.........:" << setw(n) << stride1 << endl;
	cout << ">> stride2.........:" << setw(n) << stride2 << endl;
	cout << ">>                  " << endl;
	cout << ">> nLeft...........:" << setw(n) << nLeft << endl;
	cout << ">> nRight..........:" << setw(n) << nRight << endl;
	cout << ">> nTop............:" << setw(n) << nTop << endl;
	cout << ">> nBottom.........:" << setw(n) << nBottom << endl;
	cout << ">> paddedRows......:" << setw(n) << paddedRows << endl;
	cout << ">> paddedCols......:" << setw(n) << paddedCols << endl;
	cout << ">> nOutputRows.....:" << setw(n) << nOutputRows << endl;
	cout << ">> nOutputCols.....:" << setw(n) << nOutputCols << endl;
	cout << ">>                  " << endl;
	cout << ">> prepareMultMap..:" << setw(n) << prepareMultMap.rows() 
								   << setw(n) << prepareMultMap.cols() << endl;
	cout << ">> OutputMaps......:" << setw(n) << OutputMaps.rows() 
								   << setw(n) << OutputMaps.cols() << endl;
}

void PiaPool::PrintResult(const std::vector<MatrixXd> &dst)
{

	for (int i = 0; i < (int)dst.size(); i++) {
//		cout << ">> ................:"
		cout << ">>              dst:" << i << endl;
		cout << dst[i] << endl;
	}
}

#endif
